package com.zjh.webShell.handle;


import com.jcraft.jsch.ChannelShell;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.zjh.webShell.utils.JsonUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;

@Component
@Slf4j
public class WebSSHHandler extends TextWebSocketHandler {

    private static final Map<WebSocketSession, Session> sessionMap = new ConcurrentHashMap<>();
    private static final Map<WebSocketSession, ChannelShell> channelMap = new ConcurrentHashMap<>();

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        Map<String, String> msg = JsonUtil.parseJson(payload);
        if (msg.isEmpty()) {
            log.info("msg is Empty");
            return;
        }
        String operate = msg.get("operate");
        if (StringUtils.isBlank(operate)) {
            log.info("operate is Empty");
            return;
        }
        switch (operate) {
            case "connect":
                connectToSSH(session, msg);
                break;
            case "command":
                sendCommand(session, msg.get("command"));
                break;
            case "resize":
                resizeTerminal(session, Integer.parseInt(msg.get("cols")), Integer.parseInt(msg.get("rows")));
                break;
        }
    }

    private void connectToSSH(WebSocketSession session, Map<String, String> config) {
        try {
            JSch jSch = new JSch();
            Session sshSession = jSch.getSession(config.get("username"), config.get("host"),
                    Integer.parseInt(config.getOrDefault("port", "22")));

            sshSession.setPassword(config.get("password"));
            Properties configProps = new Properties();
            configProps.put("StrictHostKeyChecking", "no");
            sshSession.setConfig(configProps);
            sshSession.connect();

            ChannelShell channel = (ChannelShell) sshSession.openChannel("shell");
            channel.setPtyType("xterm");
            channel.setPtySize(80, 50, 640, 480);

            InputStream inputFromChannel = channel.getInputStream();
            OutputStream outputToChannel = channel.getOutputStream();

            channel.connect();

            sessionMap.put(session, sshSession);
            channelMap.put(session, channel);

            // 启动线程读取SSH输出
            new Thread(() -> {
                try {
                    byte[] buffer = new byte[1024];
                    int i;
                    while ((i = inputFromChannel.read(buffer)) != -1) {
                        session.sendMessage(new TextMessage(Arrays.copyOf(buffer, i)));
                    }
                } catch (IOException e) {
                    try {
                        session.sendMessage(new TextMessage("SSH连接异常关闭: " + e.getMessage()));
                        session.close();
                    } catch (IOException ex) {
                        ex.printStackTrace();
                    }
                } finally {
                    closeSSH(session);
                }
            }).start();

        } catch (JSchException | IOException e) {
            try {
                session.sendMessage(new TextMessage("SSH连接失败: " + e.getMessage()));
                session.close();
            } catch (IOException ex) {
                ex.printStackTrace();
            }
        }
    }

    private void sendCommand(WebSocketSession session, String command) throws IOException {
        ChannelShell channel = channelMap.get(session);
        if (channel != null && channel.isConnected()) {
            OutputStream outputToChannel = channel.getOutputStream();
            outputToChannel.write(command.getBytes());
            outputToChannel.flush();
        }
    }

    private void resizeTerminal(WebSocketSession session, int cols, int rows) throws JSchException {
        ChannelShell channel = channelMap.get(session);
        if (channel != null && channel.isConnected()) {
            channel.setPtySize(cols, rows, cols * 8, rows * 16);
        }
    }

    private void closeSSH(WebSocketSession session) {
        ChannelShell channel = channelMap.get(session);
        Session sshSession = sessionMap.get(session);

        if (channel != null) {
            channel.disconnect();
            channelMap.remove(session);
        }

        if (sshSession != null) {
            sshSession.disconnect();
            sessionMap.remove(session);
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) throws Exception {
        closeSSH(session);
        super.afterConnectionClosed(session, status);
    }
}